{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8b50e216-6915-4644-b4c7-434a6fa2f843",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet llama-index-core llama-index-utils-workflow llama-index-llms-openai llama-index-graph-stores-neo4j "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5e873d91-0953-483c-ad27-922986b36ad0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/tomazbratanic/anaconda3/lib/python3.11/site-packages/pandas/core/arrays/masked.py:60: UserWarning: Pandas requires version '1.3.6' or newer of 'bottleneck' (version '1.3.5' currently installed).\n",
      "  from pandas.core import (\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.llms.anthropic import Anthropic\n",
    "from neo4j.exceptions import CypherSyntaxError\n",
    "\n",
    "from llama_index.core import ChatPromptTemplate\n",
    "from pydantic import BaseModel, Field\n",
    "from typing import List, Literal, Union, Optional\n",
    "\n",
    "from llama_index.core.workflow import (\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Workflow,\n",
    "    step,\n",
    "    Event,\n",
    "    Context,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7012b2a4-de4f-4a64-9945-8f68a79870d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore\n",
    "\n",
    "graph_store = Neo4jPropertyGraphStore(\n",
    "    username=\"recommendations\",\n",
    "    password=\"recommendations\",\n",
    "    database=\"recommendations\",\n",
    "    url=\"neo4j+s://demo.neo4jlabs.com:7687\",\n",
    "    enhanced_schema=True,\n",
    "    create_indexes=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "59248381-356a-41b5-b580-dc0bba48db00",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"ANTHROPIC_API_KEY\"] = \"sk-ant-api03-\"\n",
    "llm = Anthropic(model=\"claude-3-5-sonnet-latest\", max_tokens=8076)\n",
    "fast_llm = Anthropic(model=\"claude-3-5-haiku-20241022\", max_tokens=8076)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2cf99343-8559-4f98-ab8b-355febf59649",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Guardrail(BaseModel):\n",
    "    \"\"\"Guardrail\"\"\"\n",
    "\n",
    "    decision: Literal[\"movie\", \"end\"] = Field(\n",
    "        description=\"Decision on whether the question is related to movies\"\n",
    "        \n",
    "    )\n",
    "guardrails_system_prompt = \"\"\"You are helpful assistant\"\"\"\n",
    "guardrails_user_prompt = \"\"\"\n",
    "You are an intelligent assistant specializing in identifying movie-related questions. Your task is to determine whether a given question is related to movies or not.\n",
    "\n",
    "Here's the question you need to analyze:\n",
    "\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\n",
    "Instructions:\n",
    "1. Carefully read and understand the question provided.\n",
    "2. Determine if the question is related to movies. This includes topics such as:\n",
    "   - Specific films\n",
    "   - Actors or actresses\n",
    "   - Directors or producers\n",
    "   - Film industry topics\n",
    "   - Movie genres\n",
    "   - Film history\n",
    "   - Cinema technology\n",
    "   - Any other movie-related subjects\n",
    "\n",
    "3. In <analysis> tags, break down your thought process as follows:\n",
    "   a. List key words or phrases from the question that might indicate it's movie-related.\n",
    "   b. Consider arguments for why the question could be movie-related.\n",
    "   c. Consider arguments for why the question might not be movie-related.\n",
    "   d. Make a final decision based on the strength of these arguments.\n",
    "\n",
    "4. After your analysis, provide your final decision:\n",
    "   - If the question is related to movies, output exactly: movie\n",
    "   - If the question is not related to movies, output exactly: end\n",
    "\"\"\"\n",
    "# Refine Prompt\n",
    "chat_refine_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        guardrails_system_prompt,\n",
    "    ),\n",
    "    (\"user\", \"The question is: {question}\"),\n",
    "]\n",
    "guardrails_template = ChatPromptTemplate.from_messages(chat_refine_msgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "e618e42c-6ee1-434e-be80-986b5bfe607f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SubqueriesOutput(BaseModel):\n",
    "    \"\"\"Defines the output format for transforming a question into parallel-optimized retrieval steps.\"\"\"\n",
    "    question_breakdown: str = Field(description=(\"\"\"Before creating the final query plan, please go through a thorough analysis phase. Wrap your analysis inside <question_breakdown> tags to show your thought process and break down the question. In your analysis:\n",
    "1. Identify key entities and relationships mentioned in the question.\n",
    "2. List potential Cypher patterns that might be useful for addressing the question.\n",
    "3. Consider whether the question can be answered with a single Cypher statement or requires multiple steps.\n",
    "4. If multiple steps are needed, consider potential optimization strategies, such as indexing or query rewriting.\n",
    "5. Estimate the complexity of each proposed step (e.g., O(n), O(n^2), etc.).\n",
    "6. Identify any opportunities for parallelization.\n",
    "7. Identify potential constraints or limitations in the graph structure that might affect the query.\n",
    "8. Consider different query execution paths and their trade-offs in terms of performance and resource usage.\n",
    "9. Note any assumptions you're making about the graph structure or data availability.\n",
    "10. Keep it concise enough\"\"\"))\n",
    "    query_plan: List[List[str]] = Field(description=(\"\"\"A list of query groups where:\n",
    "        - Each group (inner list) contains queries that can be executed in parallel\n",
    "        - Groups are ordered by dependency (earlier groups must be executed before later ones)\n",
    "        - Each query must be a specific information retrieval request\"\"\"))\n",
    "\n",
    "subqueries_system = \"\"\"You are a helpful assistant\"\"\"\n",
    "\n",
    "subqueries_user = \"\"\"You are an advanced Query Planning Optimizer for Graph Databases, specifically designed to work with Cypher queries. Your task is to analyze complex questions about graph data and create efficient, structured query plans that can be translated into Cypher statements.\n",
    "\n",
    "Here is the question you need to analyze and create a query plan for:\n",
    "\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\n",
    "Please follow these steps to create an optimized query plan:\n",
    "\n",
    "1. Analyze the Question:\n",
    "Begin by thoroughly analyzing the question. Wrap your analysis in <question_analysis> tags. In your analysis:\n",
    "a. Restate the question in your own words to ensure understanding.\n",
    "b. List and number each entity and relationship mentioned in the question.\n",
    "c. For each entity and relationship, note potential Cypher patterns and estimate their complexity.\n",
    "d. Explicitly state whether the question requires a single Cypher statement or multiple steps.\n",
    "e. For multi-step queries, list potential parallelization opportunities.\n",
    "f. Enumerate any potential constraints or limitations in the graph structure.\n",
    "g. Brainstorm potential edge cases that might affect the query results.\n",
    "\n",
    "2. Create an Optimized Query Plan:\n",
    "After your analysis, create an optimized query plan. Wrap your plan development process in <plan_development> tags. Your plan should:\n",
    "- Utilize Cypher's strengths in handling relationships, aggregations, and filtering.\n",
    "- Minimize sequential dependencies while maintaining logical correctness.\n",
    "- Organize independent queries into parallel groups when possible.\n",
    "- Use a single query for tasks that don't require breakdown into smaller steps.\n",
    "- Ensure each step directly contributes to answering the user's question or preparing data for subsequent steps.\n",
    "- Express steps as high-level queries or operations, not as specific Cypher statements.\n",
    "- Address potential edge cases or limitations identified in the analysis.\n",
    "- If a simple Cypher statement can address the entire question, implement it as a single operation in the query plan.\n",
    "- For each step, list the input data required and the output data produced.\n",
    "- Explicitly state any assumptions made about the graph structure.\n",
    "- Note any potential performance bottlenecks and how they're addressed.\n",
    "\n",
    "3. Present the Query Plan:\n",
    "Present your final query plan in <query_plan> tags. The query plan should be a structured text format where:\n",
    "- Each sequential step is numbered and clearly labeled with a descriptive name.\n",
    "- Parallel operations within a step are indicated with appropriate indentation and markers.\n",
    "- Dependencies between steps are clearly indicated.\n",
    "- The order of execution is explicit and easy to follow.\n",
    "- Step and operation names describe the information being retrieved or processed.\n",
    "- Any considerations for edge cases or limitations are included in relevant steps.\n",
    "\n",
    "Here's an example of the expected query plan structure:\n",
    "\n",
    "Note: If the entire query can be efficiently executed in a single Cypher statement, present it as a single step with one operation.\n",
    "\n",
    "<query_plan>\n",
    "Step 1: Identify Tom Hanks' Most Frequent Coactor\n",
    "  Purpose: Find the actor who has appeared most often with Tom Hanks\n",
    "  Parallel: false\n",
    "  Operations:\n",
    "    - Operation 1.1: Match Tom Hanks in the database\n",
    "    - Operation 1.2: Find all coactors and count shared movies\n",
    "    - Operation 1.3: Select coactor with highest count\n",
    "    - Operation 1.4: Store coactor information for next step\n",
    "\n",
    "Step 2: Count Movies for Both Actors\n",
    "  Purpose: Count total movies for both Leonardo DiCaprio and Tom Hanks' most frequent coactor\n",
    "  Parallel: true\n",
    "  Operations:\n",
    "    - Operation 2.1: Count distinct movies for Leonardo DiCaprio\n",
    "    - Operation 2.2: Count distinct movies for identified coactor from Step 1\n",
    "\n",
    "Step 3: Compare Results\n",
    "  Purpose: Compare movie counts and determine who made more movies\n",
    "  Parallel: false\n",
    "  Operations:\n",
    "    - Operation 3.1: Compare movie counts\n",
    "    - Operation 3.2: Format result with actor names and their respective movie counts\n",
    "\n",
    "Notes:\n",
    "- Edge Case Handling: If multiple coactors tie for most frequent, use the first one alphabetically\n",
    "- Null Handling: Return appropriate message if Tom Hanks has no coactors\n",
    "- Performance: Use indexes on Person labels and name properties if available\n",
    "- Distinct Movies: Ensure counting unique movies only\n",
    "</query_plan>\n",
    "\n",
    "Please proceed with your analysis and query plan for the given question.\"\"\"\n",
    "\n",
    "query_decompose_msgs = [\n",
    "    (\"system\", subqueries_system),\n",
    "    (\"user\", subqueries_user)\n",
    "]\n",
    "\n",
    "subquery_template = ChatPromptTemplate.from_messages(query_decompose_msgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "86aa49dd-f20e-48c9-9b0b-bd910077b5b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import re\n",
    "\n",
    "def extract_query_plan(text, original_question):\n",
    "    # Handle list input by joining with newlines\n",
    "    if isinstance(text, list):\n",
    "        text = '\\n'.join(text)\n",
    "    elif not isinstance(text, str):\n",
    "        raise TypeError(\"Input must be either a string or a list of strings\")\n",
    "    \n",
    "    # Find the content between <query_plan> tags\n",
    "    plan_match = re.search(r'<query_plan>(.*?)</query_plan>', text, re.DOTALL)\n",
    "    if not plan_match:\n",
    "        return None\n",
    "    \n",
    "    plan_text = plan_match.group(1)\n",
    "    \n",
    "    # Initialize the result structure\n",
    "    query_plan = {\n",
    "        \"steps\": []\n",
    "    }\n",
    "    \n",
    "    # Split into steps\n",
    "    step_pattern = r'Step \\d+: (.*?)(?=Step \\d+:|$)'\n",
    "    steps = re.finditer(step_pattern, plan_text, re.DOTALL)\n",
    "    \n",
    "    for step_match in steps:\n",
    "        step_content = step_match.group(0)\n",
    "        \n",
    "        # Extract step number\n",
    "        step_num = re.search(r'Step (\\d+):', step_content)\n",
    "        step_num = int(step_num.group(1)) if step_num else 0\n",
    "        \n",
    "        # Extract step title\n",
    "        step_title = re.search(r'Step \\d+: (.*?)\\n', step_content)\n",
    "        step_title = step_title.group(1).strip() if step_title else \"\"\n",
    "        \n",
    "        # Extract purpose\n",
    "        purpose = re.search(r'Purpose: (.*?)\\n', step_content)\n",
    "        purpose = purpose.group(1).strip() if purpose else \"\"\n",
    "        \n",
    "        # Extract parallel setting\n",
    "        parallel = re.search(r'Parallel: (.*?)\\n', step_content)\n",
    "        parallel = parallel.group(1).strip().lower() == 'true' if parallel else None\n",
    "        \n",
    "        # Extract operations\n",
    "        operations = []\n",
    "        ops = re.finditer(r'Operation \\d+\\.\\d+: (.*?)\\n', step_content)\n",
    "        for op in ops:\n",
    "            operations.append(op.group(1).strip())\n",
    "            \n",
    "        # Extract output\n",
    "        output = re.search(r'Output: (.*?)\\n', step_content)\n",
    "        output = output.group(1).strip() if output else \"\"\n",
    "        \n",
    "        # Create step dictionary\n",
    "        step = {\n",
    "            \"step_number\": step_num,\n",
    "            \"title\": step_title,\n",
    "        }\n",
    "        \n",
    "        # Only add non-empty values\n",
    "        if purpose:\n",
    "            step[\"purpose\"] = purpose\n",
    "        if parallel is not None:\n",
    "            step[\"parallel\"] = parallel\n",
    "        if operations:\n",
    "            step[\"operations\"] = operations\n",
    "        if output:\n",
    "            step[\"output\"] = output\n",
    "            \n",
    "        query_plan[\"steps\"].append(step)\n",
    "    \n",
    "    # Extract notes if present\n",
    "    notes_match = re.search(r'Notes:(.*?)(?=</query_plan>|$)', plan_text, re.DOTALL)\n",
    "    if notes_match:\n",
    "        notes_text = notes_match.group(1)\n",
    "        notes = {}\n",
    "        \n",
    "        # Extract different note categories\n",
    "        categories = re.finditer(r'- (.*?):(.*?)(?=-|$)', notes_text, re.DOTALL)\n",
    "        for category in categories:\n",
    "            category_name = category.group(1).strip()\n",
    "            category_content = category.group(2).strip()\n",
    "            \n",
    "            # Convert bullet points to list\n",
    "            items = [item.strip('* ').strip() for item in category_content.split('\\n') if item.strip()]\n",
    "            if items:\n",
    "                notes[category_name] = items\n",
    "                \n",
    "        if notes:\n",
    "            query_plan[\"notes\"] = notes\n",
    "    # Turn into lists of lists\n",
    "    plan_steps = []\n",
    "\n",
    "    # Iterate through the steps in the plan\n",
    "    for step in query_plan['steps']:\n",
    "        step_operations = []\n",
    "\n",
    "        # If operations are parallel, each becomes an individual list item\n",
    "        if step['parallel']:\n",
    "            for operation in step['operations']:\n",
    "                #reasoning = f\"Reasoning: {step['purpose']}\"\n",
    "                step_operations.append(f\"The original question was: {original_question}. Here, you have to resolve: {operation}\")\n",
    "\n",
    "        # If not parallel, concatenate all operations with reasoning\n",
    "        else:\n",
    "            joined_operations = f\"The original question was: {original_question}. Here, you have to resolve: {step['purpose']}. Follow these steps:\" + \", \".join(\n",
    "                [f\"{operation}\" for operation in step['operations']]\n",
    "            )\n",
    "            step_operations.append(joined_operations)\n",
    "\n",
    "        # Add the operations for this step to the result list\n",
    "        plan_steps.append(step_operations)\n",
    "\n",
    "    return plan_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "578f1b93-be9a-4d51-b9b6-6bf5d54fcca7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['The original question was: Which movie has the biggest cast?. Here, you have to resolve: Identify the movie with the most distinct actors. Follow these steps:Match all Movie-Actor relationships, Group by movie, Count distinct actors per movie, Order by actor count descending, Select top result, Return movie details with cast size']]\n"
     ]
    }
   ],
   "source": [
    "#print(extract_query_plan(llm.chat(subquery_template.format_messages(question=\"Which movie has the biggest cast?\")).message.content, \"Which movie has the biggest cast?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "205b660b-b32c-4e39-9824-7b0daf4fa221",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor??. Here, you have to resolve: Find the actor who has appeared most often with Tom Hanks. Follow these steps:Match Tom Hanks node, Find all coactors through shared movies, Count shared movies per coactor, Select coactor with highest count (use ORDER BY for ties)'], ['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor??. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio', 'The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor??. Here, you have to resolve: Count distinct movies for identified coactor from Step 1'], ['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor??. Here, you have to resolve: Determine who made more movies and format output. Follow these steps:Compare movie counts, Format result with actor names and counts']]\n"
     ]
    }
   ],
   "source": [
    "#print(extract_query_plan(llm.chat(subquery_template.format_messages(question=\"Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor??\")).message.content,\"Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor??\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c4907938-fa16-4dca-b068-23a80793e68f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[['The original question was: List directors who have directed at least two movies starring the same actor. What are those movies?. Here, you have to resolve: Identify directors who have worked with the same actor in multiple movies. Follow these steps:Match director-movie-actor patterns, Group relationships by director and actor, Filter for minimum movie count, Collect movie information']]\n"
     ]
    }
   ],
   "source": [
    "#print(extract_query_plan(llm.chat(subquery_template.format_messages(question=\"List directors who have directed at least two movies starring the same actor. What are those movies?\")).message.content, \"List directors who have directed at least two movies starring the same actor. What are those movies?\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cbf05240-d619-499c-83fe-dd1f917c55dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def guardrails_step(question):\n",
    "    guardrails_output = (\n",
    "        fast_llm.as_structured_llm(Guardrail)\n",
    "        .complete(guardrails_template.format(question=question))\n",
    "        .raw\n",
    "    ).decision\n",
    "    if guardrails_output == 'end':\n",
    "        context = \"The question is not about movies or their case, so I cannot answer this question\"\n",
    "        return {\"next_event\": \"generate_final_answer\", \"arguments\": {\"context\": context, \"question\": question}}\n",
    "    # Refactor into separate step\n",
    "    queries_output = llm.chat(subquery_template.format_messages(question=question)).message.content\n",
    "    \n",
    "    return {\"next_event\": \"generate_cypher\", \"arguments\": {\"plan\": extract_query_plan(queries_output, question), \"question\": question}}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8314df30-cf3c-4441-92c1-ae88f9d0a672",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'next_event': 'generate_cypher',\n",
       " 'arguments': {'plan': [[\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the highest-rated movie Tom Hanks appeared in. Follow these steps:Match Tom Hanks to his movies, Order by rating/acclaim metric, Select top-rated movie\"],\n",
       "   [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the director of the movie from Step 1. Follow these steps:Match director to selected movie, Store director information\"],\n",
       "   [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify actor with most collaborations with selected director. Follow these steps:Match all movies by director, Match all actors in these movies, Count collaborations per actor, Select actor with highest count\"],\n",
       "   [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count Leonardo DiCaprio's movies\",\n",
       "    \"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count identified collaborator's movies\"],\n",
       "   [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Determine who has more movies. Follow these steps:Compare movie counts, Format result with actor names and counts\"]],\n",
       "  'question': \"Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??\"}}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "guardrails_step(\"Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fd6844fa-f5ad-42c2-8312-43541410b7e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.schema import TextNode\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "embed_model = OpenAIEmbedding(model=\"text-embedding-3-small\")\n",
    "\n",
    "\n",
    "examples = [\n",
    "    {\n",
    "        \"question\": \"How many artists are there?\",\n",
    "        \"query\": \"MATCH (a:Person)-[:ACTED_IN]->(:Movie) RETURN count(DISTINCT a)\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Which actors played in the movie Casino?\",\n",
    "        \"query\": \"MATCH (m:Movie {title: 'Casino'})<-[:ACTED_IN]-(a) RETURN a.name\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How many movies has Tom Hanks acted in?\",\n",
    "        \"query\": \"MATCH (a:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie) RETURN count(m)\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"List all the genres of the movie Schindler's List\",\n",
    "        \"query\": \"MATCH (m:Movie {title: 'Schindler's List'})-[:IN_GENRE]->(g:Genre) RETURN g.name\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Which actors have worked in movies from both the comedy and action genres?\",\n",
    "        \"query\": \"MATCH (a:Person)-[:ACTED_IN]->(:Movie)-[:IN_GENRE]->(g1:Genre), (a)-[:ACTED_IN]->(:Movie)-[:IN_GENRE]->(g2:Genre) WHERE g1.name = 'Comedy' AND g2.name = 'Action' RETURN DISTINCT a.name\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Which directors have made movies with at least three different actors named 'John'?\",\n",
    "        \"query\": \"MATCH (d:Person)-[:DIRECTED]->(m:Movie)<-[:ACTED_IN]-(a:Person) WHERE a.name STARTS WITH 'John' WITH d, COUNT(DISTINCT a) AS JohnsCount WHERE JohnsCount >= 3 RETURN d.name\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Identify movies where directors also played a role in the film.\",\n",
    "        \"query\": \"MATCH (p:Person)-[:DIRECTED]->(m:Movie), (p)-[:ACTED_IN]->(m) RETURN m.title, p.name\",\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Find the actor with the highest number of movies in the database.\",\n",
    "        \"query\": \"MATCH (a:Actor)-[:ACTED_IN]->(m:Movie) RETURN a.name, COUNT(m) AS movieCount ORDER BY movieCount DESC LIMIT 1\",\n",
    "    },\n",
    "]\n",
    "\"\"\"\n",
    "few_shot_nodes = []\n",
    "for line in examples:\n",
    "    few_shot_nodes.append(TextNode(text=f\"{{'query':{line['query']}, 'question': {line['question']}))\"))\n",
    "\n",
    "few_shot_index = VectorStoreIndex(few_shot_nodes, embed_model=embed_model)\n",
    "few_shot_retriever = few_shot_index.as_retriever(similarity_top_k=5)\n",
    "\"\"\"\n",
    "\n",
    "def get_fewshots(question):\n",
    "    return [] #[el.text for el in few_shot_retriever.retrieve(question)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c0597674-4e7d-4723-8094-3f1adeee3ab2",
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_system = \"\"\"Given an input question, convert it to a Cypher query. No pre-amble.\n",
    "Do not wrap the response in any backticks or anything else. Respond with a Cypher statement only!\"\"\"\n",
    "\n",
    "generate_user = \"\"\"You are a Neo4j expert. Given an input question, create a syntactically correct Cypher query to run.\n",
    "Do not wrap the response in any backticks or anything else. Respond with a Cypher statement only!\n",
    "Here is the schema information\n",
    "{schema}\n",
    "\n",
    "Below are a number of examples of questions and their corresponding Cypher queries.\n",
    "\n",
    "{fewshot_examples}\n",
    "\n",
    "User input: {question}\n",
    "Cypher query:\"\"\"\n",
    "\n",
    "generate_cypher_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        generate_system,\n",
    "    ),\n",
    "    (\"user\", generate_user),\n",
    "]\n",
    "\n",
    "text2cypher_prompt = ChatPromptTemplate.from_messages(generate_cypher_msgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4ec54fde-6995-4823-9150-b9798763d504",
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = graph_store.get_schema_str(exclude_types=[\"Actor\", \"Director\"])\n",
    "\n",
    "async def generate_cypher(subquery):\n",
    "    fewshot_examples = get_fewshots(subquery)\n",
    "    resp = await llm.achat(text2cypher_prompt.format_messages(question=subquery, schema=schema, fewshot_examples=fewshot_examples))\n",
    "    return resp.message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "7804a532-64a0-41dc-830b-c50bf75d75fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "validate_cypher_system = \"\"\"You are a helpful assistant\"\"\"\n",
    "\n",
    "validate_cypher_user = \"\"\"You are a specialized parser designed to analyze Cypher query statements and extract node property filters. Your task is to identify and extract properties used in WHERE clauses and pattern matching conditions, but only when they contain explicit literal values.\n",
    "\n",
    "Here is the Cypher statement you need to analyze:\n",
    "\n",
    "<cypher_statement>\n",
    "{cypher}\n",
    "</cypher_statement>\n",
    "\n",
    "Please follow these steps to analyze the Cypher statement and extract the relevant filters:\n",
    "\n",
    "1. Carefully read through the Cypher statement.\n",
    "2. Identify all node labels and their associated property filters.\n",
    "3. For each property filter:\n",
    "   a. Determine if it matches against a literal value (string, number, or boolean).\n",
    "   b. If it does, extract the property key and its matching literal value.\n",
    "   c. If it doesn't (e.g., property-to-property comparisons), ignore it.\n",
    "4. Format the extracted filters as a JSON object containing a \"filters\" array.\n",
    "\n",
    "Rules for extraction:\n",
    "- Only extract filters that match against literal values:\n",
    "  * Quoted strings (e.g., 'John', \"London\")\n",
    "  * Numbers (e.g., 30, 42.5)\n",
    "  * Booleans (true, false)\n",
    "- Include property filters from both MATCH patterns (e.g., (p:Person {name: 'John'})) and WHERE clauses with literal values (e.g., WHERE p.age > 30)\n",
    "- Handle both simple equality and comparison operators with literal values\n",
    "- Ignore property-to-property comparisons (e.g., WHERE m1.rating = m2.rating)\n",
    "- Ignore variable references or dynamic values\n",
    "\n",
    "Before providing the final output, break down the Cypher statement in <cypher_breakdown> tags:\n",
    "\n",
    "1. List all node labels and their associated properties from the MATCH clauses.\n",
    "2. Identify and list all WHERE clause conditions.\n",
    "3. For each property filter, explicitly state whether it matches a literal value and why.\n",
    "4. Count the number of valid filters identified.\n",
    "\n",
    "This breakdown will help ensure a thorough interpretation of the data.\n",
    "\n",
    "Output Format:\n",
    "After your analysis, provide the final output as a JSON object with a \"filters\" array. Each element in the array should be an object with the following properties:\n",
    "- node_label: The label of the node (e.g., \"Person\")\n",
    "- property_key: The property name being filtered (e.g., \"age\")\n",
    "- property_value: The literal value being matched (e.g., 30)\n",
    "\n",
    "Example of the desired output structure:\n",
    "{\n",
    "  \"filters\": [\n",
    "    {\n",
    "      \"node_label\": \"ExampleNode\",\n",
    "      \"property_key\": \"exampleProperty\",\n",
    "      \"property_value\": \"exampleValue\"\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "\n",
    "Please proceed with your analysis and provide the extracted filters in the specified JSON format.\"\"\"\n",
    "\n",
    "validate_cypher_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        validate_cypher_system,\n",
    "    ),\n",
    "    (\"user\", validate_cypher_user),\n",
    "]\n",
    "\n",
    "validate_cypher_prompt = ChatPromptTemplate.from_messages(validate_cypher_msgs)\n",
    "\n",
    "class Property(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents a filter condition based on a specific node property in a graph in a Cypher statement.\n",
    "    \"\"\"\n",
    "\n",
    "    node_label: str = Field(\n",
    "        description=\"The label of the node to which this property belongs.\"\n",
    "    )\n",
    "    property_key: str = Field(description=\"The key of the property being filtered.\")\n",
    "    property_value: str = Field(\n",
    "        description=\"The value that the property is being matched against.\"\n",
    "    )\n",
    "\n",
    "\n",
    "class ValidateCypherOutput(BaseModel):\n",
    "    \"\"\"\n",
    "    Represents the applied filters of a Cypher query's output.\n",
    "    \"\"\"\n",
    "    filters: Optional[List[Property]] = Field(\n",
    "        description=\"A list of property-based filters applied in the Cypher statement.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7795a143-f6e5-4198-bd9d-3efc1d98a036",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.graph_stores.neo4j import CypherQueryCorrector, Schema\n",
    "\n",
    "# Cypher query corrector is experimental\n",
    "corrector_schema = [\n",
    "    Schema(el[\"start\"], el[\"type\"], el[\"end\"])\n",
    "    for el in graph_store.get_schema().get(\"relationships\")\n",
    "]\n",
    "cypher_query_corrector = CypherQueryCorrector(corrector_schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ba0a4ea0-63be-4309-b5ce-75dd9229a4bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def validate_cypher(question, cypher):\n",
    "    \"\"\"\n",
    "    Validates the Cypher statements and maps any property values to the database.\n",
    "    \"\"\"\n",
    "    errors = []\n",
    "    mapping_errors = []\n",
    "    # Check for syntax errors\n",
    "    try:\n",
    "        graph_store.structured_query(f\"EXPLAIN {cypher}\")\n",
    "    except CypherSyntaxError as e:\n",
    "        errors.append(e.message)\n",
    "    # Experimental feature for correcting relationship directions\n",
    "    corrected_cypher = cypher_query_corrector(cypher)\n",
    "    if not corrected_cypher:\n",
    "        errors.append(\"The generated Cypher statement doesn't fit the graph schema\")\n",
    "    # Use LLM for mapping for values\n",
    "    llm_output =   (\n",
    "        fast_llm.as_structured_llm(ValidateCypherOutput)\n",
    "        .complete(validate_cypher_prompt.format(cypher=cypher))\n",
    "        .raw\n",
    "    )\n",
    "    if llm_output.filters:\n",
    "        for filter in llm_output.filters:\n",
    "            # Do mapping only for string values\n",
    "            try:\n",
    "                if (\n",
    "                    not [\n",
    "                        prop\n",
    "                        for prop in graph_store.get_schema()[\"node_props\"][\n",
    "                            filter.node_label\n",
    "                        ]\n",
    "                        if prop[\"property\"] == filter.property_key\n",
    "                    ][0][\"type\"]\n",
    "                    == \"STRING\"\n",
    "                ):\n",
    "                    continue\n",
    "            except: # if property is hallucinated/doesn't exist in the schema # ToDo handle it better\n",
    "                continue\n",
    "            mapping = graph_store.structured_query(\n",
    "                f\"MATCH (n:{filter.node_label}) WHERE toLower(n.`{filter.property_key}`) = toLower($value) RETURN 'yes' LIMIT 1\",\n",
    "                {\"value\": filter.property_value},\n",
    "            )\n",
    "            if not mapping:\n",
    "                mapping_errors.append(\n",
    "                    f\"Could not find node in graph with label '{filter.node_label}' where property '{filter.property_key}' equals '{filter.property_value}'. \"\n",
    "                    f\"Without this information, I cannot provide a complete answer to your question. \"\n",
    "                    f\"If you meant something else, please rephrase your question or verify the specific {filter.property_key} you're asking about. \"\n",
    "                    f\"Would you like to try with a different {filter.property_key} value?\"\n",
    "                )\n",
    "    if mapping_errors:\n",
    "        next_action = \"end\"\n",
    "    elif errors:\n",
    "        next_action = \"correct_cypher\"\n",
    "    else:\n",
    "        next_action = \"execute_cypher\"\n",
    "\n",
    "    return {\n",
    "        \"next_action\": next_action,\n",
    "        \"cypher_statement\": corrected_cypher,\n",
    "        \"cypher_errors\": errors,\n",
    "        \"mapping_errors\": mapping_errors,\n",
    "        \"steps\": [\"validate_cypher\"],\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "44d23e86-bbba-42c4-97f7-412a6a93ac25",
   "metadata": {},
   "outputs": [],
   "source": [
    "correct_cypher_system = \"\"\"You are a helpful assistant\"\"\"\n",
    "\n",
    "correct_cypher_user = \"\"\"You are an expert Cypher developer tasked with reviewing and correcting a Cypher statement written by a junior developer. Your goal is to provide an accurate and efficient Cypher statement that addresses all the identified errors while adhering to the given schema and answering the provided question.\n",
    "\n",
    "Here is the database schema you should use as a reference:\n",
    "\n",
    "<schema>\n",
    "{schema}\n",
    "</schema>\n",
    "\n",
    "The question or task for which the Cypher statement was written:\n",
    "\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\n",
    "The original Cypher statement written by the junior developer:\n",
    "\n",
    "<original_cypher>\n",
    "{cypher}\n",
    "</original_cypher>\n",
    "\n",
    "The errors identified in the original statement:\n",
    "\n",
    "<errors>\n",
    "{errors}\n",
    "</errors>\n",
    "\n",
    "Your task is to analyze the original Cypher statement, identify the issues based on the provided errors, and create a corrected version that adheres to the schema and answers the question effectively.\n",
    "\n",
    "Please follow these steps:\n",
    "\n",
    "1. Analyze the original Cypher statement and the provided errors.\n",
    "2. Identify the specific issues that need to be addressed.\n",
    "3. Develop a corrected Cypher statement that resolves all identified errors.\n",
    "4. Ensure the corrected statement adheres to the given schema.\n",
    "5. Verify that the corrected statement effectively answers the provided question.\n",
    "\n",
    "Conduct your review inside <cypher_review> tags. In your review:\n",
    "- List out each error and its corresponding correction\n",
    "- Explain how each correction addresses the schema and question\n",
    "- Consider any potential side effects of the corrections\n",
    "\n",
    "After your review, provide the corrected Cypher statement without any additional explanations, comments, or formatting. Do not use backticks or any other wrapping characters.\n",
    "\n",
    "Remember:\n",
    "- Only output a Cypher statement as your final answer.\n",
    "- Do not include any explanations or apologies in your final output.\n",
    "- Do not respond to any questions that ask for anything other than constructing a Cypher statement.\n",
    "\n",
    "Begin your response with your review in <cypher_review> tags, followed by the corrected Cypher statement inside <cypher_statement>.\"\"\"\n",
    "\n",
    "# Correct cypher\n",
    "correct_cypher_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        correct_cypher_system,\n",
    "    ),\n",
    "    (\"user\", correct_cypher_user),\n",
    "]\n",
    "\n",
    "correct_cypher_prompt = ChatPromptTemplate.from_messages(correct_cypher_msgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "979ee298-d2a1-4134-8f9a-c3ac1b41ca92",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_correct_cypher(text):\n",
    "    # Look for content between <cypher_statement> and </cypher> (optional)\n",
    "    start_tag = \"<cypher_statement>\"\n",
    "    end_tag = \"</cypher_statement>\"\n",
    "    \n",
    "    # Find the start position\n",
    "    start_pos = text.find(start_tag)\n",
    "    if start_pos == -1:\n",
    "        return None\n",
    "    \n",
    "    # Adjust start position to exclude the tag\n",
    "    start_pos += len(start_tag)\n",
    "    \n",
    "    # Find the end position\n",
    "    end_pos = text.find(end_tag)\n",
    "    if end_pos == -1:\n",
    "        # If no end tag, take until the end of the string\n",
    "        cypher = text[start_pos:].strip()\n",
    "    else:\n",
    "        cypher = text[start_pos:end_pos].strip()\n",
    "    \n",
    "    return cypher"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "9dfe2255-8146-4082-8ada-78b5ff95a79c",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def correct_cypher(subquery, cypher, errors):\n",
    "    resp = await llm.achat(correct_cypher_prompt.format_messages(question=subquery, schema=schema, errors=errors, cypher=cypher))\n",
    "    return extract_correct_cypher(resp.message.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "dfb09de3-0f7b-4e97-a39d-80bdb5eafb2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MATCH (p:Person)\n",
      "WHERE p.name = \"Tom Hanks\"\n",
      "OPTIONAL MATCH (p)-[r:ACTED_IN]->(m:Movie)\n",
      "RETURN p.name, p.bio, p.born, p.died, p.bornIn, COLLECT(m.title) as movies\n"
     ]
    }
   ],
   "source": [
    "#print(await correct_cypher(\"Who is Tom Hanks?\", \"MATCH (a:Actor) RETURN a\", [\"The query doesn't match tom hanks specifically!\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1b2562aa-64ec-42c8-b1f4-dfa285c26751",
   "metadata": {},
   "outputs": [],
   "source": [
    "information_check_system = \"\"\"You are a helpful assistant\"\"\"\n",
    "\n",
    "information_check_user = \"\"\"\n",
    "You are an advanced Query Planning Optimizer for Graph Databases, specifically designed to work with Cypher queries. Your task is to analyze complex questions about graph data, evaluate available information, and create efficient, structured query plans that can be translated into Cypher statements.\n",
    "\n",
    "Here is the information you need to work with:\n",
    "\n",
    "<subqueries>\n",
    "{subqueries}\n",
    "</subqueries>\n",
    "\n",
    "<dynamic_notebook>\n",
    "{dynamic_notebook}\n",
    "</dynamic_notebook>\n",
    "\n",
    "<current_plan>\n",
    "{plan}\n",
    "</current_plan>\n",
    "\n",
    "<original_question>\n",
    "{question}\n",
    "</original_question>\n",
    "\n",
    "Please follow these steps to evaluate the available information and create an optimized query plan:\n",
    "\n",
    "1. Analyze the Original Question:\n",
    "Begin by thoroughly analyzing the question. In your analysis:\n",
    "a. Restate the question in your own words to ensure understanding.\n",
    "b. List and number each entity and relationship mentioned in the question.\n",
    "c. For each entity and relationship, note potential Cypher patterns and estimate their complexity.\n",
    "d. Explicitly state whether the question requires a single Cypher statement or multiple steps.\n",
    "e. For multi-step queries, list potential parallelization opportunities.\n",
    "f. Enumerate any potential constraints or limitations in the graph structure.\n",
    "g. Brainstorm potential edge cases that might affect the query results.\n",
    "\n",
    "2. Review Available Information:\n",
    "Examine the subqueries, their results, and the provided condensed information in the dynamic notebook. Assess if they collectively address all components of the question.\n",
    "\n",
    "3. Identify Information Gaps:\n",
    "Compare the requirements from the question against the available information. Highlight any missing details or incomplete data that must be retrieved to form a complete answer.\n",
    "\n",
    "4. Update and Refine the Dynamic Notebook:\n",
    "Treat the condensed information as a central knowledge base. Continuously update it with key details from subquery results and integrate new data to close gaps and establish connections between facts.\n",
    "The knowledge should be concise and stored as a text that is understandable to any human\n",
    "\n",
    "5. Create or Modify the Query Plan:\n",
    "Based on your analysis, create or modify the query plan. Your plan should:\n",
    "- Utilize Cypher's strengths in handling relationships, aggregations, and filtering.\n",
    "- Minimize sequential dependencies while maintaining logical correctness.\n",
    "- Organize independent queries into parallel groups when possible.\n",
    "- Use a single query for tasks that don't require breakdown into smaller steps.\n",
    "- Ensure each step directly contributes to answering the user's question or preparing data for subsequent steps.\n",
    "- Express steps as high-level queries or operations, not as specific Cypher statements.\n",
    "- Address potential edge cases or limitations identified in the analysis.\n",
    "- For each step, list the input data required and the output data produced.\n",
    "- Explicitly state any assumptions made about the graph structure.\n",
    "- Note any potential performance bottlenecks and how they're addressed.\n",
    "\n",
    "Wrap your analysis and plan in the following tags:\n",
    "\n",
    "<question_breakdown>\n",
    "[Your detailed analysis of the question, including:\n",
    "1. Restatement of the question\n",
    "2. Numbered list of entities and relationships\n",
    "3. Cypher patterns for each entity/relationship\n",
    "4. Single or multi-step query requirement\n",
    "5. Parallelization opportunities\n",
    "6. Potential constraints or limitations\n",
    "7. Potential edge cases]\n",
    "</question_breakdown>\n",
    "\n",
    "<information_review>\n",
    "[Your review of available information and identified gaps]\n",
    "</information_review>\n",
    "\n",
    "<dynamic_notebook>\n",
    "[Updated version of the dynamic notebook, including old and new information and connections]\n",
    "</dynamic_notebook>\n",
    "\n",
    "<query_plan>\n",
    "[Your optimized query plan, following this structure:\n",
    "\n",
    "Step 1: [Step Name]\n",
    "  Purpose: [Brief description of the step's purpose]\n",
    "  Parallel: [true/false]\n",
    "  Operations:\n",
    "    - Operation 1.1: [Description of operation]\n",
    "    - Operation 1.2: [Description of operation]\n",
    "    [...]\n",
    "\n",
    "Step 2: [Step Name]\n",
    "  [...]\n",
    "\n",
    "Notes:\n",
    "- [Any additional notes, edge case handling, assumptions, or performance considerations]\n",
    "]\n",
    "</query_plan>\n",
    "\n",
    "If the task is unsolvable due to critical missing information:\n",
    "<unsolvable_task>\n",
    "[Explanation of why the task cannot be completed, including specific missing data]\n",
    "</unsolvable_task>\n",
    "\n",
    "Remember:\n",
    "- Focus only on information retrieval and avoid reasoning/analysis tasks in the query plan.\n",
    "- Optimize for parallel execution whenever possible.\n",
    "- Maintain sequential order only when necessary due to data dependencies.\n",
    "- Centralize all knowledge in the dynamic notebook.\n",
    "- If a simple Cypher statement can address the entire question, implement it as a single operation in the query plan.\n",
    "\n",
    "Please proceed with your analysis and query plan for the given question and available information. \n",
    "\"\"\"\n",
    "\n",
    "information_check_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        information_check_system,\n",
    "    ),\n",
    "    (\"user\", information_check_user),\n",
    "]\n",
    "\n",
    "information_check_prompt = ChatPromptTemplate.from_messages(information_check_msgs)\n",
    "\n",
    "class IFOutput(BaseModel):\n",
    "    \"\"\"\n",
    "    Contains a condensed summary of the available information or additional subqueries needed to answer the original question.\n",
    "    \"\"\"\n",
    "\n",
    "    dynamic_notebook: str = Field(\n",
    "        description=\"Updated version of the dynamic notebook, including old and new information and connections\"\n",
    "    )\n",
    "    query_plan: Optional[List[List[str]]] = Field(\n",
    "        description=\"Modified version of the remaining plan steps. Each group contains queries that can be executed in parallel. Null if no remaining plan exists, all gaps have been addressed, or the task is unsolvable due to missing critical information.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "15837b79-6416-425d-8eb3-7e1d91af267c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_subqueries_for_prompt(information_checks: list) -> str:\n",
    "    \"\"\"\n",
    "    Converts a list of InformationCheck objects into a string that can be added to a prompt.\n",
    "    \n",
    "    Args:\n",
    "        information_checks (List[InformationCheck]): List of information checks to process.\n",
    "    \n",
    "    Returns:\n",
    "        str: A formatted string representing subqueries and their results.\n",
    "    \"\"\"\n",
    "    subqueries_and_results = []\n",
    "    \n",
    "    for check in information_checks:\n",
    "        # Extract the first result if available, otherwise use \"No result available.\"\n",
    "        result = (\n",
    "            check.database_output[0] if check.database_output else \"No result available.\"\n",
    "        )\n",
    "        subqueries_and_results.append(\n",
    "            f\"- Subquery: {check.subquery}\\n  Result: {result}\"\n",
    "        )\n",
    "    \n",
    "    return \"\\n\".join(subqueries_and_results)\n",
    "\n",
    "def information_check(subquery_events, original_question, dynamic_notebook, plan):\n",
    "    subqueries = format_subqueries_for_prompt(subquery_events)\n",
    "    print(f\"Before: {dynamic_notebook}\")\n",
    "    print(f\"Plan: {plan}\")\n",
    "    llm_output =   (\n",
    "        llm.as_structured_llm(IFOutput)\n",
    "        .complete(information_check_prompt.format(subqueries=subqueries, original_question=original_question, dynamic_notebook=dynamic_notebook, plan=plan), tool_choice={'type': 'tool', 'name':'IFOutput'})\n",
    "        .raw\n",
    "    )\n",
    "    print(f\"After: {llm_output.dynamic_notebook}\")\n",
    "    print(f\"New Plan: {llm_output.query_plan}\")\n",
    "    return {'dynamic_notebook': llm_output.dynamic_notebook, 'modified_plan': llm_output.query_plan}\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ed70406c-c95b-4e29-a863-cc29c14a89c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "final_answer_system = \"\"\"You are a helpful assistant.\"\"\"\n",
    "\n",
    "final_answer_user = \"\"\"\n",
    "You are a highly capable AI assistant specializing in providing accurate and concise answers based on given information. Your task is to analyze a provided context and answer a specific question using only the information contained within that context.\n",
    "\n",
    "Here is the context you will be working with:\n",
    "\n",
    "<context>\n",
    "{context}\n",
    "</context>\n",
    "\n",
    "And here is the question you need to answer:\n",
    "\n",
    "<question>\n",
    "{question}\n",
    "</question>\n",
    "\n",
    "Please follow these steps to complete your task:\n",
    "\n",
    "1. Carefully read and analyze the context and question.\n",
    "2. Wrap your thought process inside <thought_process> tags:\n",
    "   - Quote key information from the context that is relevant to the question.\n",
    "   - Consider whether the question is clear or if clarification is needed.\n",
    "   - List potential answers to the question based on the context.\n",
    "   - For each potential answer, provide arguments supporting it and evaluate the reliability and sufficiency of the information.\n",
    "   - Determine if the context provides sufficient information to answer the question confidently.\n",
    "   - If information is missing, specify what additional details are required.\n",
    "\n",
    "It's OK for this section to be quite long.\n",
    "\n",
    "3. Based on your analysis, do the following:\n",
    "   a. Ensure your answer is clear, relevant, and directly addresses the user's question. Explain your reasoning within this section.\n",
    "\n",
    "Remember:\n",
    "- Focus solely on the provided context. Do not use external knowledge or make assumptions unless explicitly stated in the context.\n",
    "- If the question is ambiguous, ask for clarification before proceeding.\n",
    "- Be concise yet thorough in your responses.\n",
    "- Maintain a professional and helpful tone throughout your response.\n",
    "\n",
    "Please begin your thought process now.\n",
    "\"\"\"\n",
    "\n",
    "final_answer_msgs = [\n",
    "    (\n",
    "        \"system\",\n",
    "        final_answer_system,\n",
    "    ),\n",
    "    (\"user\", final_answer_user),\n",
    "]\n",
    "\n",
    "final_answer_prompt = ChatPromptTemplate.from_messages(final_answer_msgs)\n",
    "\n",
    "async def generate_final_answer(question, context):\n",
    "    resp = await fast_llm.achat(final_answer_prompt.format_messages(question=question, context=context))\n",
    "    return resp.message.content                 \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "0ee5da9c-679f-4dd5-b99f-2d63334fd8c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GenerateCypher(Event):\n",
    "    subquery: str\n",
    "    \n",
    "class ValidateCypher(Event):\n",
    "    subquery: str\n",
    "    generated_cypher: str\n",
    "\n",
    "class CorrectCypher(Event):\n",
    "    cypher: str\n",
    "    subquery: str\n",
    "    errors: List[str]\n",
    "\n",
    "class ExecuteCypher(Event):\n",
    "    validated_cypher: str\n",
    "    subquery: str\n",
    "\n",
    "class InformationCheck(Event):\n",
    "    cypher: str\n",
    "    subquery: str\n",
    "    database_output: list\n",
    "    \n",
    "class GenerateFinalAnswer(Event):\n",
    "    context: str\n",
    "\n",
    "class ConcurrentFlow(Workflow):\n",
    "    @step\n",
    "    async def start(self, ctx: Context, ev: StartEvent) -> GenerateCypher | GenerateFinalAnswer:\n",
    "        original_question = ev.input\n",
    "        await ctx.set(\"original_question\", original_question)\n",
    "        await ctx.set(\"dynamic_notebook\", \"\")\n",
    "        await ctx.set(\"subqueries_cypher_history\", {})\n",
    "        guardrails_output = guardrails_step(original_question)\n",
    "        if guardrails_output.get(\"next_event\") == \"generate_final_answer\":\n",
    "            context = \"The question is not about movies or cast, so I cannot answer the question\"\n",
    "            return GenerateFinalAnswer(context=context)\n",
    "\n",
    "        # store in global context \n",
    "        subqueries = guardrails_output[\"arguments\"].get(\"plan\")\n",
    "        await ctx.set(\"count_of_subqueries\", len(subqueries[0])) #we use this in ctx.collect()\n",
    "        await ctx.set(\"plan\", subqueries[1:]) #we use this in information check\n",
    "        # Send events\n",
    "        for subquery in subqueries[0]:\n",
    "            print(subquery)\n",
    "            ctx.send_event(GenerateCypher(subquery=subquery))\n",
    "\n",
    "    @step(num_workers=4)\n",
    "    async def generate_cypher_step(self, ctx: Context, ev: GenerateCypher) -> ValidateCypher:\n",
    "        print(\"Running generate_cypher \", ev.subquery)\n",
    "        generated_cypher = await generate_cypher(ev.subquery)\n",
    "        return ValidateCypher(subquery=ev.subquery, generated_cypher=generated_cypher)\n",
    "\n",
    "    @step(num_workers=4)\n",
    "    async def validate_cypher_step(self, ctx: Context, ev: ValidateCypher) -> GenerateFinalAnswer | ExecuteCypher | CorrectCypher:\n",
    "        print(\"Running validate_cypher \", ev)\n",
    "        results = validate_cypher(ev.subquery, ev.generated_cypher)\n",
    "        if results['next_action'] == \"end\": # DB value mapping\n",
    "            return GenerateFinalAnswer(context=str(results[\"mapping_errors\"]))\n",
    "        if results['next_action'] == \"execute_cypher\":\n",
    "            return ExecuteCypher(subquery=ev.subquery, validated_cypher=ev.generated_cypher)\n",
    "        if results['next_action'] == \"correct_cypher\":\n",
    "            return CorrectCypher(subquery=ev.subquery, cypher=ev.generated_cypher, errors=results['cypher_errors'])\n",
    "\n",
    "    @step(num_workers=4)\n",
    "    async def correct_cypher_step(self, ctx: Context, ev: CorrectCypher) -> ValidateCypher:\n",
    "        print(\"Running validate_cypher \", ev)\n",
    "        results = await correct_cypher(ev.subquery, ev.cypher, ev.errors)\n",
    "        return ValidateCypher(subquery=ev.subquery, generated_cypher=results)\n",
    "    \n",
    "    @step\n",
    "    async def execute_cypher_step(self, ctx: Context, ev: ExecuteCypher) -> InformationCheck:\n",
    "        # wait until we receive all events\n",
    "        print(\"Running execute_cypher_step \", ev)\n",
    "        try:\n",
    "            database_output = graph_store.structured_query(ev.validated_cypher)\n",
    "        except Exception as e: # Dividing by zero, etc...\n",
    "            database_output = e.message\n",
    "        return InformationCheck(subquery=ev.subquery, cypher=ev.validated_cypher, database_output=database_output)\n",
    "\n",
    "    @step\n",
    "    async def information_check_step(self, ctx: Context, ev: InformationCheck) -> GenerateCypher | GenerateFinalAnswer:\n",
    "        # wait until we receive all events\n",
    "        print(\"Running information_check_step\", ev)\n",
    "        # retrieve from context\n",
    "        number_of_subqueries = await ctx.get(\"count_of_subqueries\")\n",
    "        result = ctx.collect_events(ev, [InformationCheck] * number_of_subqueries)\n",
    "        if result is None:\n",
    "            return None\n",
    "        # Add executed cypher statements to global state\n",
    "        subqueries_cypher_history = await ctx.get(\"subqueries_cypher_history\")\n",
    "        new_subqueries_cypher = {\n",
    "                item.subquery: {\n",
    "                    \"cypher\": item.cypher,\n",
    "                    \"database_output\": item.database_output\n",
    "                } for item in result\n",
    "            }\n",
    "        await ctx.set(\"subqueries_cypher_history\", {**subqueries_cypher_history, **new_subqueries_cypher})\n",
    "\n",
    "        original_question = await ctx.get(\"original_question\")\n",
    "        dynamic_notebook = await ctx.get(\"dynamic_notebook\")\n",
    "        plan = await ctx.get(\"plan\")\n",
    "\n",
    "        # Do the information check\n",
    "        \n",
    "        data = information_check(result, original_question, dynamic_notebook, plan)\n",
    "        # Go fetch additional information if needed\n",
    "        if data.get(\"modified_plan\"):\n",
    "            await ctx.set(\"count_of_subqueries\", len(data['modified_plan'][0])) # this is used for ctx.collect()\n",
    "            await ctx.set(\"dynamic_notebook\", data[\"dynamic_notebook\"])\n",
    "            await ctx.set(\"plan\", data.get(\"modified_plan\")[1:])\n",
    "            for subquery in data[\"modified_plan\"][0]:\n",
    "                ctx.send_event(GenerateCypher(subquery=subquery))\n",
    "        else:\n",
    "            return GenerateFinalAnswer(context=data['dynamic_notebook'])\n",
    "\n",
    "    @step\n",
    "    async def final_answer(self, ctx: Context, ev: GenerateFinalAnswer) -> StopEvent:\n",
    "        original_question = await ctx.get(\"original_question\")\n",
    "        subqueries_cypher_history = await ctx.get(\"subqueries_cypher_history\")\n",
    "        # wait until we receive all events\n",
    "        print(\"Running final_answer \", ev)\n",
    "        resp = await generate_final_answer(original_question, ev.context)\n",
    "        return StopEvent(result={\"text\":resp, \"subqueries_cypher_history\": subqueries_cypher_history})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "01ba2699-c0e4-4518-a0d6-3d7bd40113b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running step start\n",
      "The original question was: Who made more movies, Leonardo DiCaprio or Tom Hanks?. Here, you have to resolve: Retrieve and compare movie counts for both actors. Follow these steps:Match both actors and their movies, Count distinct movies for each actor, Compare counts and determine who made more movies, Format results with actor names and counts\n",
      "Step start produced no event\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  The original question was: Who made more movies, Leonardo DiCaprio or Tom Hanks?. Here, you have to resolve: Retrieve and compare movie counts for both actors. Follow these steps:Match both actors and their movies, Count distinct movies for each actor, Compare counts and determine who made more movies, Format results with actor names and counts\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery='The original question was: Who made more movies, Leonardo DiCaprio or Tom Hanks?. Here, you have to resolve: Retrieve and compare movie counts for both actors. Follow these steps:Match both actors and their movies, Count distinct movies for each actor, Compare counts and determine who made more movies, Format results with actor names and counts' generated_cypher=\"MATCH (leo:Person {name: 'Leonardo DiCaprio'})-[:ACTED_IN]->(leoMovies:Movie)\\nWITH leo, COUNT(DISTINCT leoMovies) as leoCount\\nMATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(tomMovies:Movie)\\nWITH leo, leoCount, tom, COUNT(DISTINCT tomMovies) as tomCount\\nRETURN \\n    leo.name as Actor1,\\n    leoCount as Actor1Movies,\\n    tom.name as Actor2,\\n    tomCount as Actor2Movies,\\n    CASE \\n        WHEN leoCount > tomCount THEN leo.name + ' made more movies'\\n        WHEN tomCount > leoCount THEN tom.name + ' made more movies'\\n        ELSE 'Both made the same number of movies'\\n    END as Comparison\"\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (leo:Person {name: 'Leonardo DiCaprio'})-[:ACTED_IN]->(leoMovies:Movie)\\nWITH leo, COUNT(DISTINCT leoMovies) as leoCount\\nMATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(tomMovies:Movie)\\nWITH leo, leoCount, tom, COUNT(DISTINCT tomMovies) as tomCount\\nRETURN \\n    leo.name as Actor1,\\n    leoCount as Actor1Movies,\\n    tom.name as Actor2,\\n    tomCount as Actor2Movies,\\n    CASE \\n        WHEN leoCount > tomCount THEN leo.name + ' made more movies'\\n        WHEN tomCount > leoCount THEN tom.name + ' made more movies'\\n        ELSE 'Both made the same number of movies'\\n    END as Comparison\" subquery='The original question was: Who made more movies, Leonardo DiCaprio or Tom Hanks?. Here, you have to resolve: Retrieve and compare movie counts for both actors. Follow these steps:Match both actors and their movies, Count distinct movies for each actor, Compare counts and determine who made more movies, Format results with actor names and counts'\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher=\"MATCH (leo:Person {name: 'Leonardo DiCaprio'})-[:ACTED_IN]->(leoMovies:Movie)\\nWITH leo, COUNT(DISTINCT leoMovies) as leoCount\\nMATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(tomMovies:Movie)\\nWITH leo, leoCount, tom, COUNT(DISTINCT tomMovies) as tomCount\\nRETURN \\n    leo.name as Actor1,\\n    leoCount as Actor1Movies,\\n    tom.name as Actor2,\\n    tomCount as Actor2Movies,\\n    CASE \\n        WHEN leoCount > tomCount THEN leo.name + ' made more movies'\\n        WHEN tomCount > leoCount THEN tom.name + ' made more movies'\\n        ELSE 'Both made the same number of movies'\\n    END as Comparison\" subquery='The original question was: Who made more movies, Leonardo DiCaprio or Tom Hanks?. Here, you have to resolve: Retrieve and compare movie counts for both actors. Follow these steps:Match both actors and their movies, Count distinct movies for each actor, Compare counts and determine who made more movies, Format results with actor names and counts' database_output=[{'Actor1': 'Leonardo DiCaprio', 'Actor1Movies': 24, 'Actor2': 'Tom Hanks', 'Actor2Movies': 38, 'Comparison': 'Tom Hanks made more movies'}]\n",
      "Before: \n",
      "Plan: []\n",
      "After: Actor Comparison:\n",
      "- Leonardo DiCaprio has appeared in 24 movies\n",
      "- Tom Hanks has appeared in 38 movies\n",
      "- Tom Hanks has made more movies than Leonardo DiCaprio (38 vs 24)\n",
      "New Plan: None\n",
      "Step information_check_step produced event GenerateFinalAnswer\n",
      "Running step final_answer\n",
      "Running final_answer  context='Actor Comparison:\\n- Leonardo DiCaprio has appeared in 24 movies\\n- Tom Hanks has appeared in 38 movies\\n- Tom Hanks has made more movies than Leonardo DiCaprio (38 vs 24)'\n",
      "Step final_answer produced event StopEvent\n",
      "{'text': '<thought_process>\\nKey Information:\\n- Leonardo DiCaprio has appeared in 24 movies\\n- Tom Hanks has appeared in 38 movies\\n- The context explicitly states that Tom Hanks has made more movies than Leonardo DiCaprio\\n\\nAnalysis:\\n1. The question directly asks who made more movies between Leonardo DiCaprio and Tom Hanks.\\n2. The context provides clear numerical evidence:\\n   - DiCaprio: 24 movies\\n   - Hanks: 38 movies\\n3. The context also directly confirms that Tom Hanks has made more movies\\n4. The information appears complete and unambiguous\\n5. No additional clarification is needed\\n\\nReliability Assessment:\\n- The information is precise and quantitative\\n- The comparison is straightforward and clear\\n- The source seems credible as it provides specific movie count details\\n\\nPotential Answers:\\n- Tom Hanks has made more movies than Leonardo DiCaprio\\n- Tom Hanks (38 movies) has a higher movie count compared to Leonardo DiCaprio (24 movies)\\n\\nConfidence Level: High - The context provides direct, unequivocal information to answer the question.\\n</thought_process>\\n\\nAnswer:\\nTom Hanks has made more movies than Leonardo DiCaprio. Specifically, Tom Hanks has appeared in 38 movies, while Leonardo DiCaprio has appeared in 24 movies.', 'subqueries_cypher_history': {'The original question was: Who made more movies, Leonardo DiCaprio or Tom Hanks?. Here, you have to resolve: Retrieve and compare movie counts for both actors. Follow these steps:Match both actors and their movies, Count distinct movies for each actor, Compare counts and determine who made more movies, Format results with actor names and counts': {'cypher': \"MATCH (leo:Person {name: 'Leonardo DiCaprio'})-[:ACTED_IN]->(leoMovies:Movie)\\nWITH leo, COUNT(DISTINCT leoMovies) as leoCount\\nMATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(tomMovies:Movie)\\nWITH leo, leoCount, tom, COUNT(DISTINCT tomMovies) as tomCount\\nRETURN \\n    leo.name as Actor1,\\n    leoCount as Actor1Movies,\\n    tom.name as Actor2,\\n    tomCount as Actor2Movies,\\n    CASE \\n        WHEN leoCount > tomCount THEN leo.name + ' made more movies'\\n        WHEN tomCount > leoCount THEN tom.name + ' made more movies'\\n        ELSE 'Both made the same number of movies'\\n    END as Comparison\", 'database_output': [{'Actor1': 'Leonardo DiCaprio', 'Actor1Movies': 24, 'Actor2': 'Tom Hanks', 'Actor2Movies': 38, 'Comparison': 'Tom Hanks made more movies'}]}}}\n",
      "Took 0:00:40.966543\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "start = datetime.now()\n",
    "w = ConcurrentFlow(timeout=120, verbose=True)\n",
    "result = await w.run(input=\"Who made more movies, Leonardo DiCaprio or Tom Hanks?\")\n",
    "print(result)\n",
    "print(f\"Took {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "97362343-ff0d-40d2-a679-c487dbcb58b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Find the actor who has appeared in the most movies with Tom Hanks. Follow these steps:Match Tom Hanks and all coactors in shared movies, Count shared movies per coactor, Select coactor with highest count (use name as tiebreaker)\n",
      "Running generate_cypher  The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Find the actor who has appeared in the most movies with Tom Hanks. Follow these steps:Match Tom Hanks and all coactors in shared movies, Count shared movies per coactor, Select coactor with highest count (use name as tiebreaker)\n",
      "Running validate_cypher  subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Find the actor who has appeared in the most movies with Tom Hanks. Follow these steps:Match Tom Hanks and all coactors in shared movies, Count shared movies per coactor, Select coactor with highest count (use name as tiebreaker)' generated_cypher=\"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWHERE tom <> coactor\\nWITH coactor, COUNT(m) as movies, coactor.name as name\\nORDER BY movies DESC, name ASC\\nLIMIT 1\\nRETURN coactor.name as MostFrequentCostar, movies as SharedMovies\"\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWHERE tom <> coactor\\nWITH coactor, COUNT(m) as movies, coactor.name as name\\nORDER BY movies DESC, name ASC\\nLIMIT 1\\nRETURN coactor.name as MostFrequentCostar, movies as SharedMovies\" subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Find the actor who has appeared in the most movies with Tom Hanks. Follow these steps:Match Tom Hanks and all coactors in shared movies, Count shared movies per coactor, Select coactor with highest count (use name as tiebreaker)'\n",
      "Running information_check_step cypher=\"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWHERE tom <> coactor\\nWITH coactor, COUNT(m) as movies, coactor.name as name\\nORDER BY movies DESC, name ASC\\nLIMIT 1\\nRETURN coactor.name as MostFrequentCostar, movies as SharedMovies\" subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Find the actor who has appeared in the most movies with Tom Hanks. Follow these steps:Match Tom Hanks and all coactors in shared movies, Count shared movies per coactor, Select coactor with highest count (use name as tiebreaker)' database_output=[{'MostFrequentCostar': 'Tim Allen', 'SharedMovies': 4}]\n",
      "Before: \n",
      "Plan: [['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio', 'The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1'], ['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Determine who made more movies and format output. Follow these steps:Compare movie counts, Format result with actor names and counts']]\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' most frequent coactor is Tim Allen, appearing together in 4 movies\n",
      "2. Need to determine: Number of movies for Leonardo DiCaprio\n",
      "3. Need to determine: Number of movies for Tim Allen\n",
      "4. Final comparison needed between Leonardo DiCaprio and Tim Allen's movie counts\n",
      "New Plan: [['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio', 'The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1'], ['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Determine who made more movies and format output. Follow these steps:Compare movie counts, Format result with actor names and counts']]\n",
      "Running generate_cypher  The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio\n",
      "Running generate_cypher  The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1\n",
      "Running validate_cypher  subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio' generated_cypher='MATCH (actor:Person {name: \"Leonardo DiCaprio\"})-[:ACTED_IN]->(movie:Movie)\\nRETURN actor.name, COUNT(DISTINCT movie) as movieCount'\n",
      "Running execute_cypher_step  validated_cypher='MATCH (actor:Person {name: \"Leonardo DiCaprio\"})-[:ACTED_IN]->(movie:Movie)\\nRETURN actor.name, COUNT(DISTINCT movie) as movieCount' subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio'\n",
      "Running information_check_step cypher='MATCH (actor:Person {name: \"Leonardo DiCaprio\"})-[:ACTED_IN]->(movie:Movie)\\nRETURN actor.name, COUNT(DISTINCT movie) as movieCount' subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio' database_output=[{'actor.name': 'Leonardo DiCaprio', 'movieCount': 24}]\n",
      "Running validate_cypher  subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1' generated_cypher=\"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWITH coactor, count(DISTINCT m) as moviesWithTom\\nORDER BY moviesWithTom DESC\\nLIMIT 1\\nWITH coactor\\nMATCH (coactor)-[:ACTED_IN]->(m:Movie)\\nRETURN coactor.name, count(DISTINCT m) as totalMovies\"\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWITH coactor, count(DISTINCT m) as moviesWithTom\\nORDER BY moviesWithTom DESC\\nLIMIT 1\\nWITH coactor\\nMATCH (coactor)-[:ACTED_IN]->(m:Movie)\\nRETURN coactor.name, count(DISTINCT m) as totalMovies\" subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1'\n",
      "Running information_check_step cypher=\"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWITH coactor, count(DISTINCT m) as moviesWithTom\\nORDER BY moviesWithTom DESC\\nLIMIT 1\\nWITH coactor\\nMATCH (coactor)-[:ACTED_IN]->(m:Movie)\\nRETURN coactor.name, count(DISTINCT m) as totalMovies\" subquery='The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1' database_output=[{'coactor.name': 'Tim Allen', 'totalMovies': 12}]\n",
      "Before: Key Facts:\n",
      "1. Tom Hanks' most frequent coactor is Tim Allen, appearing together in 4 movies\n",
      "2. Need to determine: Number of movies for Leonardo DiCaprio\n",
      "3. Need to determine: Number of movies for Tim Allen\n",
      "4. Final comparison needed between Leonardo DiCaprio and Tim Allen's movie counts\n",
      "Plan: [['The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Determine who made more movies and format output. Follow these steps:Compare movie counts, Format result with actor names and counts']]\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' most frequent coactor is Tim Allen, appearing together in 4 movies\n",
      "2. Leonardo DiCaprio has appeared in 24 movies\n",
      "3. Tim Allen has appeared in 12 movies\n",
      "4. Comparison Result: Leonardo DiCaprio (24 movies) has made more movies than Tim Allen (12 movies)\n",
      "New Plan: None\n",
      "Running final_answer  context=\"Key Facts:\\n1. Tom Hanks' most frequent coactor is Tim Allen, appearing together in 4 movies\\n2. Leonardo DiCaprio has appeared in 24 movies\\n3. Tim Allen has appeared in 12 movies\\n4. Comparison Result: Leonardo DiCaprio (24 movies) has made more movies than Tim Allen (12 movies)\"\n",
      "{'text': \"<thought_process>\\nKey observations from the context:\\n- Tom Hanks' most frequent coactor is Tim Allen\\n- Tim Allen has appeared in 12 movies\\n- Leonardo DiCaprio has appeared in 24 movies\\n\\nAnalysis of the question:\\n- The question asks to compare the number of movies between Leonardo DiCaprio and Tom Hanks' most frequent coactor (Tim Allen)\\n- We have clear movie count information for both individuals\\n\\nPotential answers:\\n1. Leonardo DiCaprio (24 movies) has made more movies than Tim Allen (12 movies)\\n   - This is directly supported by the context\\n   - The comparison is explicitly stated in the context\\n   - The information appears reliable and precise\\n\\nConfidence in answer:\\n- The context provides a direct, unambiguous comparison\\n- No additional information is needed to answer the question\\n- The numerical data is clear and specific\\n</thought_process>\\n\\nAnswer:\\nLeonardo DiCaprio has made more movies than Tom Hanks' most frequent coactor, Tim Allen. Specifically, DiCaprio has appeared in 24 movies, while Tim Allen has appeared in 12 movies. The context directly confirms this comparison, showing that DiCaprio has made twice as many movies as Tim Allen.\", 'subqueries_cypher_history': {'The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Find the actor who has appeared in the most movies with Tom Hanks. Follow these steps:Match Tom Hanks and all coactors in shared movies, Count shared movies per coactor, Select coactor with highest count (use name as tiebreaker)': {'cypher': \"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWHERE tom <> coactor\\nWITH coactor, COUNT(m) as movies, coactor.name as name\\nORDER BY movies DESC, name ASC\\nLIMIT 1\\nRETURN coactor.name as MostFrequentCostar, movies as SharedMovies\", 'database_output': [{'MostFrequentCostar': 'Tim Allen', 'SharedMovies': 4}]}, 'The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio': {'cypher': 'MATCH (actor:Person {name: \"Leonardo DiCaprio\"})-[:ACTED_IN]->(movie:Movie)\\nRETURN actor.name, COUNT(DISTINCT movie) as movieCount', 'database_output': [{'actor.name': 'Leonardo DiCaprio', 'movieCount': 24}]}, 'The original question was: Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?. Here, you have to resolve: Count distinct movies for identified coactor from Step 1': {'cypher': \"MATCH (tom:Person {name: 'Tom Hanks'})-[:ACTED_IN]->(m:Movie)<-[:ACTED_IN]-(coactor:Person)\\nWITH coactor, count(DISTINCT m) as moviesWithTom\\nORDER BY moviesWithTom DESC\\nLIMIT 1\\nWITH coactor\\nMATCH (coactor)-[:ACTED_IN]->(m:Movie)\\nRETURN coactor.name, count(DISTINCT m) as totalMovies\", 'database_output': [{'coactor.name': 'Tim Allen', 'totalMovies': 12}]}}}\n",
      "Took 0:01:03.070041\n"
     ]
    }
   ],
   "source": [
    "start = datetime.now()\n",
    "w = ConcurrentFlow(timeout=120, verbose=False)\n",
    "result = await w.run(input=\"Who made more movies, Leonardo di Caprio or Tom Hanks most frequent coactor?\")\n",
    "print(result)\n",
    "print(f\"Took {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "143e4ba1-0950-41fa-919e-efbbaa7d3010",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running step start\n",
      "The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the movie(s) with highest rating where Tom Hanks acted. Follow these steps:Match Tom Hanks node, Get all his movies with ratings, Sort by rating and select highest, Store movie information for next step\n",
      "Step start produced no event\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the movie(s) with highest rating where Tom Hanks acted. Follow these steps:Match Tom Hanks node, Get all his movies with ratings, Sort by rating and select highest, Store movie information for next step\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery=\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the movie(s) with highest rating where Tom Hanks acted. Follow these steps:Match Tom Hanks node, Get all his movies with ratings, Sort by rating and select highest, Store movie information for next step\" generated_cypher='MATCH (actor:Person {name: \"Tom Hanks\"})-[:ACTED_IN]->(movie:Movie)\\nWHERE movie.imdbRating IS NOT NULL\\nRETURN movie\\nORDER BY movie.imdbRating DESC\\nLIMIT 1'\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher='MATCH (actor:Person {name: \"Tom Hanks\"})-[:ACTED_IN]->(movie:Movie)\\nWHERE movie.imdbRating IS NOT NULL\\nRETURN movie\\nORDER BY movie.imdbRating DESC\\nLIMIT 1' subquery=\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the movie(s) with highest rating where Tom Hanks acted. Follow these steps:Match Tom Hanks node, Get all his movies with ratings, Sort by rating and select highest, Store movie information for next step\"\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher='MATCH (actor:Person {name: \"Tom Hanks\"})-[:ACTED_IN]->(movie:Movie)\\nWHERE movie.imdbRating IS NOT NULL\\nRETURN movie\\nORDER BY movie.imdbRating DESC\\nLIMIT 1' subquery=\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Find the movie(s) with highest rating where Tom Hanks acted. Follow these steps:Match Tom Hanks node, Get all his movies with ratings, Sort by rating and select highest, Store movie information for next step\" database_output=[{'movie': {'languages': ['English'], 'year': 1998, 'imdbId': '0120570', 'runtime': 60, 'imdbRating': 8.9, 'movieId': '27002', 'countries': ['USA'], 'imdbVotes': 7460, 'title': 'From the Earth to the Moon', 'url': 'https://themoviedb.org/movie/30983', 'tmdbId': '30983', 'plot': 'Dramatized portrayal of the Apollo manned space program.', 'released': '1998-04-05'}}]\n",
      "Before: \n",
      "Plan: [[\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify the director of the selected movie. Follow these steps:Match director relationship with stored movie, Store director information\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify actor with most collaborations with found director. Follow these steps:Match all movies by director, Find all actors in these movies, Count collaborations per actor, Select actor with highest count\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio\", \"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count distinct movies for identified collaborator\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Determine who has more movies and format response. Follow these steps:Compare movie counts, Format result with actor names and counts\"]]\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "New Plan: [[\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify the director of the selected movie. Follow these steps:Match director relationship with stored movie, Store director information\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify actor with most collaborations with found director. Follow these steps:Match all movies by director, Find all actors in these movies, Count collaborations per actor, Select actor with highest count\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio\", \"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count distinct movies for identified collaborator\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Determine who has more movies and format response. Follow these steps:Compare movie counts, Format result with actor names and counts\"]]\n",
      "Step information_check_step produced no event\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify the director of the selected movie. Follow these steps:Match director relationship with stored movie, Store director information\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery=\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify the director of the selected movie. Follow these steps:Match director relationship with stored movie, Store director information\" generated_cypher=\"MATCH (hanks:Person {name: 'Tom Hanks'})\\nMATCH (hanks)-[:ACTED_IN]->(m:Movie)\\nWITH m ORDER BY m.imdbRating DESC LIMIT 1\\nMATCH (director:Person)-[:DIRECTED]->(m)\\nWITH director\\nMATCH (director)-[:DIRECTED]->(movies:Movie)<-[:ACTED_IN]-(costar:Person)\\nWITH costar, count(movies) as collaborations\\nORDER BY collaborations DESC LIMIT 1\\nWITH costar\\nMATCH (dicaprio:Person {name: 'Leonardo DiCaprio'})\\nMATCH (dicaprio)-[:ACTED_IN]->(m1:Movie)\\nWITH costar, count(m1) as dicaprio_count\\nMATCH (costar)-[:ACTED_IN]->(m2:Movie)\\nWITH costar, dicaprio_count, count(m2) as costar_count\\nRETURN \\n    CASE \\n        WHEN dicaprio_count > costar_count THEN 'Leonardo DiCaprio'\\n        WHEN dicaprio_count < costar_count THEN costar.name\\n        ELSE 'Equal'\\n    END as winner,\\n    dicaprio_count,\\n    costar_count,\\n    costar.name as costar\"\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (hanks:Person {name: 'Tom Hanks'})\\nMATCH (hanks)-[:ACTED_IN]->(m:Movie)\\nWITH m ORDER BY m.imdbRating DESC LIMIT 1\\nMATCH (director:Person)-[:DIRECTED]->(m)\\nWITH director\\nMATCH (director)-[:DIRECTED]->(movies:Movie)<-[:ACTED_IN]-(costar:Person)\\nWITH costar, count(movies) as collaborations\\nORDER BY collaborations DESC LIMIT 1\\nWITH costar\\nMATCH (dicaprio:Person {name: 'Leonardo DiCaprio'})\\nMATCH (dicaprio)-[:ACTED_IN]->(m1:Movie)\\nWITH costar, count(m1) as dicaprio_count\\nMATCH (costar)-[:ACTED_IN]->(m2:Movie)\\nWITH costar, dicaprio_count, count(m2) as costar_count\\nRETURN \\n    CASE \\n        WHEN dicaprio_count > costar_count THEN 'Leonardo DiCaprio'\\n        WHEN dicaprio_count < costar_count THEN costar.name\\n        ELSE 'Equal'\\n    END as winner,\\n    dicaprio_count,\\n    costar_count,\\n    costar.name as costar\" subquery=\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify the director of the selected movie. Follow these steps:Match director relationship with stored movie, Store director information\"\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher=\"MATCH (hanks:Person {name: 'Tom Hanks'})\\nMATCH (hanks)-[:ACTED_IN]->(m:Movie)\\nWITH m ORDER BY m.imdbRating DESC LIMIT 1\\nMATCH (director:Person)-[:DIRECTED]->(m)\\nWITH director\\nMATCH (director)-[:DIRECTED]->(movies:Movie)<-[:ACTED_IN]-(costar:Person)\\nWITH costar, count(movies) as collaborations\\nORDER BY collaborations DESC LIMIT 1\\nWITH costar\\nMATCH (dicaprio:Person {name: 'Leonardo DiCaprio'})\\nMATCH (dicaprio)-[:ACTED_IN]->(m1:Movie)\\nWITH costar, count(m1) as dicaprio_count\\nMATCH (costar)-[:ACTED_IN]->(m2:Movie)\\nWITH costar, dicaprio_count, count(m2) as costar_count\\nRETURN \\n    CASE \\n        WHEN dicaprio_count > costar_count THEN 'Leonardo DiCaprio'\\n        WHEN dicaprio_count < costar_count THEN costar.name\\n        ELSE 'Equal'\\n    END as winner,\\n    dicaprio_count,\\n    costar_count,\\n    costar.name as costar\" subquery=\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify the director of the selected movie. Follow these steps:Match director relationship with stored movie, Store director information\" database_output=[]\n",
      "Before: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "Plan: [[\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Identify actor with most collaborations with found director. Follow these steps:Match all movies by director, Find all actors in these movies, Count collaborations per actor, Select actor with highest count\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count distinct movies for Leonardo DiCaprio\", \"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Count distinct movies for identified collaborator\"], [\"The original question was: Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??. Here, you have to resolve: Determine who has more movies and format response. Follow these steps:Compare movie counts, Format result with actor names and counts\"]]\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "2. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "New Plan: [['MATCH (m:Movie {id: 27002})<-[:DIRECTED]-(d:Person) RETURN d AS director'], ['MATCH (d:Person)-[:DIRECTED]->(m:Movie)<-[:ACTED_IN]-(a:Person) WHERE d.id = [previous_director_id] WITH a, COUNT(DISTINCT m) as collaboration_count ORDER BY collaboration_count DESC LIMIT 1 RETURN a AS most_frequent_costar, collaboration_count'], [\"MATCH (p:Person {name: 'Leonardo DiCaprio'})-[:ACTED_IN]->(m:Movie) RETURN COUNT(DISTINCT m) as movie_count\"]]\n",
      "Step information_check_step produced no event\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  MATCH (m:Movie {id: 27002})<-[:DIRECTED]-(d:Person) RETURN d AS director\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery='MATCH (m:Movie {id: 27002})<-[:DIRECTED]-(d:Person) RETURN d AS director' generated_cypher=\"MATCH (m:Movie {movieId: '27002'})<-[:DIRECTED]-(d:Person) RETURN d AS director\"\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (m:Movie {movieId: '27002'})<-[:DIRECTED]-(d:Person) RETURN d AS director\" subquery='MATCH (m:Movie {id: 27002})<-[:DIRECTED]-(d:Person) RETURN d AS director'\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher=\"MATCH (m:Movie {movieId: '27002'})<-[:DIRECTED]-(d:Person) RETURN d AS director\" subquery='MATCH (m:Movie {id: 27002})<-[:DIRECTED]-(d:Person) RETURN d AS director' database_output=[]\n",
      "Before: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "2. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "Plan: [['MATCH (d:Person)-[:DIRECTED]->(m:Movie)<-[:ACTED_IN]-(a:Person) WHERE d.id = [previous_director_id] WITH a, COUNT(DISTINCT m) as collaboration_count ORDER BY collaboration_count DESC LIMIT 1 RETURN a AS most_frequent_costar, collaboration_count'], [\"MATCH (p:Person {name: 'Leonardo DiCaprio'})-[:ACTED_IN]->(m:Movie) RETURN COUNT(DISTINCT m) as movie_count\"]]\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "2. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified (attempted query returned no results)\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "3. Query Status:\n",
      "   - Attempted to find director for Movie ID 27002 but query returned no results\n",
      "   - This suggests potential data inconsistency or missing relationships in the graph\n",
      "New Plan: [['MATCH (m:Movie {id: 27002}) RETURN m.title, m.year, m.rating', 'MATCH (m:Movie {id: 27002})<-[:ACTED_IN]-(p:Person) RETURN p.name AS actor'], ['MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre']]\n",
      "Step information_check_step produced no event\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  MATCH (m:Movie {id: 27002}) RETURN m.title, m.year, m.rating\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  MATCH (m:Movie {id: 27002})<-[:ACTED_IN]-(p:Person) RETURN p.name AS actor\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery='MATCH (m:Movie {id: 27002})<-[:ACTED_IN]-(p:Person) RETURN p.name AS actor' generated_cypher=\"MATCH (m:Movie {movieId: '27002'})<-[:ACTED_IN]-(p:Person) \\nRETURN p.name AS actor\"\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (m:Movie {movieId: '27002'})<-[:ACTED_IN]-(p:Person) \\nRETURN p.name AS actor\" subquery='MATCH (m:Movie {id: 27002})<-[:ACTED_IN]-(p:Person) RETURN p.name AS actor'\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher=\"MATCH (m:Movie {movieId: '27002'})<-[:ACTED_IN]-(p:Person) \\nRETURN p.name AS actor\" subquery='MATCH (m:Movie {id: 27002})<-[:ACTED_IN]-(p:Person) RETURN p.name AS actor' database_output=[{'actor': 'David Andrews'}, {'actor': 'Lane Smith'}, {'actor': 'Tom Hanks'}, {'actor': 'Nick Searcy'}]\n",
      "Step information_check_step produced no event\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery='MATCH (m:Movie {id: 27002}) RETURN m.title, m.year, m.rating' generated_cypher=\"MATCH (m:Movie {movieId: '27002'}) \\nRETURN m.title, m.year, m.imdbRating\"\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (m:Movie {movieId: '27002'}) \\nRETURN m.title, m.year, m.imdbRating\" subquery='MATCH (m:Movie {id: 27002}) RETURN m.title, m.year, m.rating'\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher=\"MATCH (m:Movie {movieId: '27002'}) \\nRETURN m.title, m.year, m.imdbRating\" subquery='MATCH (m:Movie {id: 27002}) RETURN m.title, m.year, m.rating' database_output=[{'m.title': 'From the Earth to the Moon', 'm.year': 1998, 'm.imdbRating': 8.9}]\n",
      "Before: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "2. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified (attempted query returned no results)\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "3. Query Status:\n",
      "   - Attempted to find director for Movie ID 27002 but query returned no results\n",
      "   - This suggests potential data inconsistency or missing relationships in the graph\n",
      "Plan: [['MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre']]\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "2. Known Cast Members:\n",
      "   - David Andrews (Actor)\n",
      "3. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified (attempted query returned no results)\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "4. Query Status:\n",
      "   - Attempted to find director for Movie ID 27002 but query returned no results\n",
      "   - This suggests potential data inconsistency or missing relationships in the graph\n",
      "   - Pending genre information query for Movie ID 27002\n",
      "New Plan: [['MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre']]\n",
      "Step information_check_step produced no event\n",
      "Running step generate_cypher_step\n",
      "Running generate_cypher  MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre\n",
      "Step generate_cypher_step produced event ValidateCypher\n",
      "Running step validate_cypher_step\n",
      "Running validate_cypher  subquery='MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre' generated_cypher=\"MATCH (m:Movie {movieId: '27002'})-[:IN_GENRE]->(g:Genre) \\nRETURN g.name AS genre\"\n",
      "Step validate_cypher_step produced event ExecuteCypher\n",
      "Running step execute_cypher_step\n",
      "Running execute_cypher_step  validated_cypher=\"MATCH (m:Movie {movieId: '27002'})-[:IN_GENRE]->(g:Genre) \\nRETURN g.name AS genre\" subquery='MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre'\n",
      "Step execute_cypher_step produced event InformationCheck\n",
      "Running step information_check_step\n",
      "Running information_check_step cypher=\"MATCH (m:Movie {movieId: '27002'})-[:IN_GENRE]->(g:Genre) \\nRETURN g.name AS genre\" subquery='MATCH (m:Movie {id: 27002})-[:IN_GENRE]->(g:Genre) RETURN g.name AS genre' database_output=[{'genre': 'Action'}, {'genre': 'Documentary'}, {'genre': 'Drama'}, {'genre': 'Thriller'}]\n",
      "Before: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "2. Known Cast Members:\n",
      "   - David Andrews (Actor)\n",
      "3. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified (attempted query returned no results)\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "4. Query Status:\n",
      "   - Attempted to find director for Movie ID 27002 but query returned no results\n",
      "   - This suggests potential data inconsistency or missing relationships in the graph\n",
      "   - Pending genre information query for Movie ID 27002\n",
      "Plan: []\n",
      "After: Key Facts:\n",
      "1. Tom Hanks' highest rated movie is \"From the Earth to the Moon\" (1998)\n",
      "   - IMDB Rating: 8.9\n",
      "   - Movie ID: 27002\n",
      "   - Plot: Dramatized portrayal of the Apollo manned space program\n",
      "   - Genre: Action\n",
      "2. Known Cast Members:\n",
      "   - David Andrews (Actor)\n",
      "3. Missing Critical Information:\n",
      "   - Director of \"From the Earth to the Moon\" not yet identified (attempted query returned no results)\n",
      "   - No information about potential co-stars with the director\n",
      "   - No information about Leonardo DiCaprio's movie count\n",
      "4. Query Status:\n",
      "   - Attempted to find director for Movie ID 27002 but query returned no results\n",
      "   - Genre information retrieved: Action\n",
      "   - This suggests potential data inconsistency or missing relationships in the graph\n",
      "New Plan: None\n",
      "Step information_check_step produced event GenerateFinalAnswer\n",
      "Running step final_answer\n",
      "Running final_answer  context='Key Facts:\\n1. Tom Hanks\\' highest rated movie is \"From the Earth to the Moon\" (1998)\\n   - IMDB Rating: 8.9\\n   - Movie ID: 27002\\n   - Plot: Dramatized portrayal of the Apollo manned space program\\n   - Genre: Action\\n2. Known Cast Members:\\n   - David Andrews (Actor)\\n3. Missing Critical Information:\\n   - Director of \"From the Earth to the Moon\" not yet identified (attempted query returned no results)\\n   - No information about potential co-stars with the director\\n   - No information about Leonardo DiCaprio\\'s movie count\\n4. Query Status:\\n   - Attempted to find director for Movie ID 27002 but query returned no results\\n   - Genre information retrieved: Action\\n   - This suggests potential data inconsistency or missing relationships in the graph'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Exception in callback Dispatcher.span.<locals>.wrapper.<locals>.handle_future_result(span_id='Workflow.run...-4034a4ee6ad1', bound_args=<BoundArgumen...ed movie??\"})>, instance=<__main__.Con...t 0x36de908d0>, context=<_contextvars...t 0x36a64ef40>)(<WorkflowHandler cancelled>) at /Users/tomazbratanic/anaconda3/lib/python3.11/site-packages/llama_index/core/instrumentation/dispatcher.py:273\n",
      "handle: <Handle Dispatcher.span.<locals>.wrapper.<locals>.handle_future_result(span_id='Workflow.run...-4034a4ee6ad1', bound_args=<BoundArgumen...ed movie??\"})>, instance=<__main__.Con...t 0x36de908d0>, context=<_contextvars...t 0x36a64ef40>)(<WorkflowHandler cancelled>) at /Users/tomazbratanic/anaconda3/lib/python3.11/site-packages/llama_index/core/instrumentation/dispatcher.py:273>\n",
      "Traceback (most recent call last):\n",
      "  File \"/Users/tomazbratanic/anaconda3/lib/python3.11/asyncio/events.py\", line 80, in _run\n",
      "    self._context.run(self._callback, *self._args)\n",
      "  File \"/Users/tomazbratanic/anaconda3/lib/python3.11/site-packages/llama_index/core/instrumentation/dispatcher.py\", line 283, in handle_future_result\n",
      "    exception = future.exception()\n",
      "                ^^^^^^^^^^^^^^^^^^\n",
      "asyncio.exceptions.CancelledError\n"
     ]
    },
    {
     "ename": "CancelledError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mCancelledError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[29], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m start \u001b[38;5;241m=\u001b[39m datetime\u001b[38;5;241m.\u001b[39mnow()\n\u001b[1;32m      2\u001b[0m w \u001b[38;5;241m=\u001b[39m ConcurrentFlow(timeout\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m150\u001b[39m, verbose\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[0;32m----> 3\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m w\u001b[38;5;241m.\u001b[39mrun(\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mWho has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m most critically acclaimed movie??\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(result)\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTook \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mdatetime\u001b[38;5;241m.\u001b[39mnow()\u001b[38;5;250m \u001b[39m\u001b[38;5;241m-\u001b[39m\u001b[38;5;250m \u001b[39mstart\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n",
      "\u001b[0;31mCancelledError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "start = datetime.now()\n",
    "w = ConcurrentFlow(timeout=150, verbose=True)\n",
    "result = await w.run(input=\"Who has appeared in more movies: Leonardo DiCaprio or the actor who has co-starred most frequently with the director of Tom Hanks' most critically acclaimed movie??\")\n",
    "print(result)\n",
    "print(f\"Took {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da7e3317-71bf-4c0a-8a97-9e46c97699d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = datetime.now()\n",
    "w = ConcurrentFlow(timeout=60, verbose=True)\n",
    "result = await w.run(input=\"How many movies did Tomaz Bratanic acted in?\")\n",
    "print(result)\n",
    "print(f\"Took {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5797cb48-f34e-418d-94bd-366920d3ecd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = datetime.now()\n",
    "w = ConcurrentFlow(timeout=30, verbose=False)\n",
    "result = await w.run(input=\"What'up esse?\")\n",
    "print(result)\n",
    "print(f\"Took {datetime.now() - start}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f455f31f-2a02-4e84-b75a-bdd951f147aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.utils.workflow import (\n",
    "    draw_all_possible_flows,\n",
    "    draw_most_recent_execution,\n",
    ")\n",
    "\n",
    "draw_most_recent_execution(w, filename=\"joke_flow_recent.html\")\n",
    "draw_all_possible_flows(w, filename=\"joke_flow_recenst.html\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "748eb2f8-076f-4e1a-bbdd-4b8525cf3df6",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16ef0c50-0e9f-4430-9626-376c4291a099",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "# URL of the Parquet file\n",
    "url = \"https://huggingface.co/datasets/tomasonjo/text2cypher_demo_test/resolve/main/test.parquet\"\n",
    "\n",
    "# Read the Parquet file directly from the URL\n",
    "df = pd.read_parquet(url)\n",
    "\n",
    "# Display the first few rows\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cca7184-c4ca-4679-bc37-a03652349c6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = ConcurrentFlow(timeout=150, verbose=False)\n",
    "answers = []\n",
    "\"\"\"\n",
    "for i, row in df.iterrows():\n",
    "    print(i)\n",
    "    question = row['question']\n",
    "    try:\n",
    "        result = await w.run(input=question)\n",
    "        print(result)\n",
    "        answers.append(result)\n",
    "    except:\n",
    "        answers.append(None)\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0736f1ff-097a-4b50-b49e-f8373a0a57bd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
